#!/usr/bin/env python3
import asyncio
from connection_manager import ConnectionManager
from typing import Dict, Optional


class ClientManager:
    def __init__(self):
        self.clients: Dict[str, ConnectionManager] = {}
        self.lock = asyncio.Lock()

    async def _create_client_internal(self, client_id: str) -> bool:
        """内部方法：创建新客户端（假设已持有锁）"""
        if client_id in self.clients:
            return True
        self.clients[client_id] = ConnectionManager(client_id)
        return True

    async def create_client(self, client_id: str) -> bool:
        """创建一个新的客户端
        参数:
        - client_id: 客户端唯一标识符
        """
        async with self.lock:
            return await self._create_client_internal(client_id)

    async def get_client_connection_manager(self, client_id: str) -> Optional[ConnectionManager]:
        """获取指定客户端的连接管理器
        参数:
        - client_id: 客户端唯一标识符
        """
        async with self.lock:
            if client_id not in self.clients:
                await self._create_client_internal(client_id)
            return self.clients[client_id]

    async def _close_client_internal(self, client_id: str) -> int:
        """内部方法：关闭客户端（假设已持有锁）"""
        if client_id not in self.clients:
            return 0
        conn_manager = self.clients[client_id]
        count = await conn_manager.close_all_connections()
        # 删除客户端
        del self.clients[client_id]
        return count

    async def close_client(self, client_id: str) -> int:
        """关闭指定客户端的所有连接
        参数:
        - client_id: 客户端ID
        返回关闭的连接数量
        """
        async with self.lock:
            return await self._close_client_internal(client_id)

    async def close_all_clients(self) -> int:
        """关闭所有客户端的所有连接
        返回关闭的连接数量
        """
        # 获取客户端ID列表
        async with self.lock:
            client_ids = list(self.clients.keys())
        
        # 并行关闭所有客户端
        tasks = []
        for client_id in client_ids:
            tasks.append(self.close_client(client_id))
        
        # 等待所有任务完成
        results = await asyncio.gather(*tasks)
        return sum(results)

    async def list_clients(self) -> Dict[str, list]:
        """列出所有客户端及其连接
        返回格式: {client_id: [connection_info]}
        """
        # 获取客户端和连接管理器的副本
        async with self.lock:
            clients_copy = list(self.clients.items())
        
        # 并行获取每个客户端的连接信息
        result = {}
        tasks = {}
        for client_id, conn_manager in clients_copy:
            tasks[client_id] = asyncio.create_task(conn_manager.list_connections())
        
        # 等待所有任务完成并收集结果
        for client_id, task in tasks.items():
            result[client_id] = await task
        
        return result